import cv2
import numpy as np
from scipy.optimize import curve_fit


def in_interval(a, l, r):
    if (a >= l) and (a <= r):
        return True
    else:
        return False


def search_light_blocks(bin_data):
    n, m = bin_data.shape
    flag_grid = np.zeros((n, m), dtype='int')
    light_blocks = []
    for i in range(n):
        for j in range(m):
            if (not flag_grid[i, j]) and bin_data[i, j]:
                queue = [(i, j)]
                l = 0
                flag_grid[i, j] = 1
                x, y = i, j
                while l < len(queue):
                    x, y = queue[l]
                    for dx in range(-1, 2):
                        for dy in range(-1, 2):
                            if (dx, dy) != (0, 0):
                                x2 = x + dx
                                y2 = y + dy
                                if in_interval(x2, 0, n - 1) and in_interval(y2, 0, m - 1):
                                    if (not flag_grid[x2, y2]) and bin_data[x2, y2]:
                                        queue.append((x2, y2))
                                        flag_grid[x2, y2] = 1
                    l += 1

                light_blocks.append(tuple(queue))
    return light_blocks


def find_ions(data):
    median = cv2.medianBlur(data, 5)
    median = np.array(median, dtype='int')
    n, m = median.shape
    median_1D = median.reshape(n * m)

    median_1D = np.sort(median_1D)

    valley, peak = median_1D[10], median_1D[-10]

    truncate_high = median_1D[np.where(median_1D > (valley + peak) / 2)]
    truncate_low = median_1D[np.where(median_1D <= (valley + peak) / 2)]

    hist_low, bins_low = np.histogram(a=truncate_low, bins=21)
    hist_high, bins_high = np.histogram(a=truncate_high, bins=21)

    # print(hist_low)
    # print(bins_low)

    max_index_low = np.argmax(hist_low)
    max_index_high = np.argmax(hist_high)
    max_low = (bins_low[max_index_low] + bins_low[max_index_low + 1]) / 2
    max_high = (bins_high[max_index_high] + bins_high[max_index_high + 1]) / 2
    threshold = (max_low + max_high) / 2
    # print("binaryzation threshold = ", threshold)
    # print(hist_high)
    # 得到threshold之后，做二值化处理(binarization)
    bin_data = np.zeros((n, m), dtype='int')
    bin_data[np.where(median > threshold)] = 1

    light_blocks = search_light_blocks(bin_data)

    # 区分噪声点与光斑的像素点threshold, 离子光斑至少有threshold个亮的像素点聚集在一块
    threshold_block = 12

    circles = []
    for block in light_blocks:
        if len(block) > threshold_block:
            # we find a circle!
            sum_xy = np.sum(np.array(block), 0)
            avrg_x = sum_xy[0] / len(block)
            avrg_y = sum_xy[1] / len(block)
            r = 0
            for x, y in block:
                if (x - avrg_x) ** 2 + (y - avrg_y) ** 2 > r ** 2:
                    r = ((x - avrg_x) ** 2 + (y - avrg_y) ** 2) ** 0.5
            if r < min(n, m) / 3:
                circles.append((avrg_x, avrg_y, r))
    # print(len(circles))
    # print(circles)
    return circles


def gaussian(xdata, A, x0, y0, sigmax, sigmay):
    return A * np.exp(-(xdata[0] - x0) ** 2 / (2 * sigmax ** 2) - (xdata[1] - y0) ** 2 / (2 * sigmay ** 2))


# Where: a = -1/(2sigmax^2), b = -1/(2sigmay^2), c = x0/(sigmax^2), d = y0/(sigmay^2)
#       e = ln(Amp) - x0^2/(2sigmax^2) - y0^2/(2sigmay^2)
def ln_gaussian(xdata, a, b, c, d, e):
    return a * xdata[0] * xdata[0] + b * xdata[1] * xdata[1] + c * xdata[0] + d * xdata[1] + e


def fit_gaussian(data, circles):
    circles_data = []
    n, m = data.shape
    for circle in circles:
        circle_points = []
        z = []
        x, y = circle[0], circle[1]
        radius = circle[2]
        lx = round(np.floor(x - radius))
        rx = round(np.ceil(x + radius))
        ly = round(np.floor(y - radius))
        ry = round(np.ceil(y + radius))
        for i in range(lx, rx):
            for j in range(ly, ry):
                if (i - x) ** 2 + (j - y) ** 2 <= radius ** 2:
                    if in_interval(i, 0, n - 1) and in_interval(j, 0, m - 1):
                        circle_points.append([i, j])
                        z.append(data[i, j])
        circle_points = np.array(circle_points).T
        z = np.array(z)
        circles_data.append((circle_points, z))

    amp_list = []
    for i in range(len(circles_data)):
        xdata = circles_data[i][0]
        z = circles_data[i][1]
        x, y, radius = circles[i]
        amp0 = data[round(x), round(y)]
        p0 = (amp0, x, y, 1, 1)
        popt, pcov = curve_fit(gaussian, xdata, z, p0=p0, maxfev=5000)
        # print(popt)
        # amp == popt[0]
        amp_list.append(popt[0])
    return amp_list


def fit_ln_gaussian(data, circles):
    '''
    fig = plt.figure()
    ax = Axes3D(fig)
    ax.plot_surface(x,y,z)
    '''

    '''
    path = 2
    if path == 1:
        lnz = np.log(z)
        p0 = (0,0,0,0,0)
        popt, pcov = curve_fit(ln_gaussian, data, lnz, p0 = p0)
        print(popt)
        a,b,c,d,e = popt

        sigma_x = (-1/(2*a)) ** 0.5
        sigma_y = (-1/(2*b)) ** 0.5
        x0 = c * sigma_x**2
        y0 = d * sigma_y**2
        lnAmp = e + c/2 + d/2
        Amp = np.exp(lnAmp)
        print(lnAmp, Amp, x0, y0, sigma_x, sigma_y)
    elif path == 2:
    '''
    pass
